package guo.fx;

import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.LinkedList;

public class H264Parser {

    public interface Callback {
        void onFrame(ByteBuffer buffer);
    }

    //核心思想是通过队列缓存startCodeLength个数据，每次读取一个字节数据，保存到队列尾部，
    // 然后判断是否满足001或0001这个要求，满足则为一帧数据直接返回，不满足则移除头部数据并写入到ByteBuffer里面，
    // 如此循环到最后，再处理残留数据即可。
    //算法可以优化，每次读取多一些数据，然后通过KMP算法查找001或0001，然后再复制。
    public static int parseH264(InputStream inputStream, Callback callback) throws IOException {
        int frameIndex = 0;
        int startCodeLength = 0;
        LinkedList<Integer> cacheList = new LinkedList<>();
        ByteBuffer byteBuffer = ByteBuffer.allocate(2 * 1024 * 1024);

        byte[] header = new byte[4];
        int result = inputStream.read(header);
        if (result < 0) return result;

        //数据不够
        if (result != 4) return -2;

        frameIndex += 4;

        int h1 = header[0];
        int h2 = header[1];
        int h3 = header[2];
        int h4 = header[3];
        if (h1 == 0 && h2 == 0 && h3 == 1) {
            startCodeLength = 3;
            cacheList.addLast(h4); //这里要缓存
        } else if (h1 == 0 && h2 == 0 && h3 == 0 && h4 == 1) {
            startCodeLength = 4;
        } else {//不符合H264文件规范
            return -3;
        }

        while(true) {
            int value = inputStream.read();
            if (value == -1) break;
            frameIndex += 1;

            cacheList.addLast(value);

            if (cacheList.size() < startCodeLength) {
                continue;
            }

            if (startCodeLength == 3) {
                processFrameOf3Startcode(cacheList, byteBuffer, callback);
            } else if (startCodeLength == 4) {
                processFrameOf4Startcode(cacheList, byteBuffer, callback);
            } else {
                break;
            }
        }

        //最后一帧数据清理
        if (cacheList.size() > 0) {
            for (int i = 0; i < cacheList.size(); i++) {
                int temp = cacheList.get(i);
                byteBuffer.put((byte) temp);
            }
            cacheList.clear();
        }

        byteBuffer.flip();
        int bufferSize = byteBuffer.limit() - byteBuffer.position();
        if (bufferSize == 0) {
            byteBuffer.clear();
        } else {
            if (callback != null) {
                callback.onFrame(byteBuffer);
            }
            byteBuffer.clear();
        }
        return 0;
    }

    //此时就是一帧数据，返回ByteBuffer
    private static void processFrameOf3Startcode(LinkedList<Integer> list, ByteBuffer byteBuffer, Callback callback) {
        int v1 = list.get(0);
        int v2 = list.get(1);
        int v3 = list.get(2);
        if (v1 == 0 && v2 == 0 && v3 == 1) {
            list.clear();
            //切换到读模式
            byteBuffer.flip();
            int bufferSize = byteBuffer.limit() - byteBuffer.position();
            if (bufferSize == 0) {
                byteBuffer.clear();
                return;
            }
            if(callback != null) {
                callback.onFrame(byteBuffer);
            }
            //清空buffer，再次读取一帧
            byteBuffer.clear();
        } else {
            int temp = list.removeFirst();
            byteBuffer.put((byte) temp);
        }
    }

    //此时就是一帧数据，返回ByteBuffer
    private static void processFrameOf4Startcode(LinkedList<Integer> list, ByteBuffer byteBuffer, Callback callback) {
        int v1 = list.get(0);
        int v2 = list.get(1);
        int v3 = list.get(2);
        int v4 = list.get(3);
        if (v1 == 0 && v2 == 0 && v3 == 0 && v4 == 1) {
            list.clear();
            //切换到读模式
            byteBuffer.flip();
            int bufferSize = byteBuffer.limit() - byteBuffer.position();
            if (bufferSize == 0) {
                byteBuffer.clear();
                return;
            }
            if(callback != null) {
                callback.onFrame(byteBuffer);
            }
            //清空buffer，再次读取一帧
            byteBuffer.clear();
        } else {
            int temp = list.removeFirst();
            byteBuffer.put((byte) temp);
        }
    }
}
